package core

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"github.com/kataras/golog"
	"github.com/thoas/go-funk"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"net"
	"ntrace/enviroments"
	"ntrace/utils/routines"
	"time"
)

type UDPReceiver struct {
	BaseReceiver
	UDPConn *ipv4.RawConn
}

func NewUDPReceiver(ctx context.Context, task *TraceTask, stat *TraceStat) *UDPReceiver {
	return &UDPReceiver{
		BaseReceiver: BaseReceiver{
			Ctx:  ctx,
			Name: "ICMPReceiver",
			Task: task,
			Stat: stat,
		},
	}
}

func (r *UDPReceiver) StartReceiver() error {
	routines.GoWithRecover("UDPReceiver-icmp", r.doReceiveICMP)
	routines.GoWithRecover("UDPReceiver-udp", r.doReceiveUDP)
	return nil
}

func (r *UDPReceiver) StopReceiver() {
	r.BaseReceiver.StopReceiver()
	if r.UDPConn != nil {
		r.UDPConn.Close()
	}

}

func (r *UDPReceiver) doReceiveICMP() {
	golog.Debugf("UDPReceiver doReceiveICMP")

	conn, err := OpenReceiverIP4PacketConn("ip4:icmp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.Conn = conn

	filter := CreateBPFUDPICMP(net.ParseIP(r.Task.DstIP))
	err = SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			conn.Close()
			return
		default:
			r.receiveOneICMPPkt(payloadBuf)
		}
	}
}

func (r *UDPReceiver) receiveOneICMPPkt(payloadBuf []byte) {
	conn := r.Conn
	err := conn.SetReadDeadline(time.Now().Add(time.Second))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		golog.Debugf("read packet failed. err: %+v", err)
		return
	}
	rxTs := time.Now().UnixNano()

	icmpMsg, err := icmp.ParseMessage(1, payload)

	devIP := h.Src.To4().String()
	var originIPPkt []byte
	var reachDst bool
	switch icmpMsg.Type {
	case ipv4.ICMPTypeDestinationUnreachable:
		originIPPkt = icmpMsg.Body.(*icmp.DstUnreach).Data
		reachDst = true
	case ipv4.ICMPTypeTimeExceeded:
		originIPPkt = icmpMsg.Body.(*icmp.TimeExceeded).Data
		reachDst = false
	default:
		return
	}

	_ = r.updateICMPRecord(originIPPkt, devIP, reachDst, rxTs)
}

func (r *UDPReceiver) doReceiveUDP() {
	golog.Debugf("UDPReceiver doReceiveUDP")

	conn, err := OpenReceiverIP4PacketConn("ip4:udp", "0.0.0.0") // 本地地址
	if err != nil {
		return
	}
	r.UDPConn = conn

	pid := enviroments.PID
	filter := CreateBPFUDP(uint16(pid))
	err = SetRawConnBPF(conn, filter)
	if err != nil {
		fmt.Println("unable to open network connection")
		return
	}

	payloadBuf := AllocRxBuf(r.Task.MsgLen)
	for {
		select {
		case <-r.Ctx.Done():
			conn.Close()
			return
		default:
			r.receiveOneUDPPkt(payloadBuf)
		}

	}
}

func (r *UDPReceiver) receiveOneUDPPkt(payloadBuf []byte) {
	conn := r.UDPConn
	err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	if err != nil {
		return
	}

	h, payload, _, err := conn.ReadFrom(payloadBuf)
	if err != nil {
		golog.Debugf("err: %+v", err)
		return
	}
	rxTs := time.Now().UnixNano()

	hLen := h.Len
	devIP := h.Src.To4().String()
	srcPort := binary.BigEndian.Uint16(payload[hLen : hLen+2])
	dstPort := binary.BigEndian.Uint16(payload[hLen+2 : hLen+4])
	ipID := binary.BigEndian.Uint16(payload[hLen+18 : hLen+20])
	if devIP == r.Task.DstIP {
		UpdateUDPPktRecord(&r.Stat.PktRecordBucket, devIP, ipID, srcPort, dstPort, devIP, rxTs, true)
	}
}

func (r *UDPReceiver) updateICMPRecord(originIPPkt []byte, devIP string, reachDst bool, rxTs int64) error {
	// is ipv4
	if len(originIPPkt) == 0 || (originIPPkt[0]>>4) != 4 {
		return errors.New("unsupported")
	}

	// ip header length
	hLen := (originIPPkt[0] & 0x0F) * 4
	srcIP := net.IPv4(originIPPkt[12], originIPPkt[13], originIPPkt[14], originIPPkt[15])
	dstIP := net.IPv4(originIPPkt[16], originIPPkt[17], originIPPkt[18], originIPPkt[19])
	// ip identifier
	ipID := binary.BigEndian.Uint16(originIPPkt[4:6])
	srcPort := binary.BigEndian.Uint16(originIPPkt[hLen : hLen+2])
	dstPort := binary.BigEndian.Uint16(originIPPkt[hLen+2 : hLen+4])
	if r.Task.DstIP == dstIP.To4().String() &&
		funk.ContainsString(r.Task.SrcIPs, srcIP.To4().String()) {

		pid := enviroments.PID
		golog.Debugf("srcIP: %s, dstIP: %s, pid: %d, deviceAddr: %s", srcIP.To4().String(), dstIP.To4().String(), pid, devIP)

		UpdateUDPPktRecord(&r.Stat.PktRecordBucket, dstIP.To4().String(), ipID, srcPort, dstPort, devIP, rxTs, reachDst)
	}

	return nil
}

func UpdateUDPPktRecord(pktRecordBucket *PktRecordBucket, target string, ipID uint16, srcPort uint16, dstPort uint16, device string, rxTs int64, reachDst bool) {
	key := fmt.Sprintf("%s-%d-%d-%d", target, ipID, srcPort, dstPort)
	golog.Debugf("receive key: %s", key)
	pktRecord := &PktRecord{
		Key:      key,
		Device:   device,
		RxTs:     rxTs,
		Received: true,
		ReachDst: reachDst,
	}

	pktRecordBucket.Update(pktRecord)
}
